import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
import torchvision
from torchvision import transforms
train_ds_original = torchvision.datasets.MNIST('data', 
                                      train=True, 
                                      transform=torchvision.transforms.ToTensor(), 
                                      download=True)

test_ds  = torchvision.datasets.MNIST('data', 
                                      train=False, 
                                      transform=torchvision.transforms.ToTensor(), 
                                      download=True)
batch_size = 32
train_dl = torch.utils.data.DataLoader(train_ds_original, 
                                       batch_size=batch_size, 
                                       shuffle=True)
train_transform = transforms.Compose([
    transforms.RandomRotation(0),      # 随机旋转±20度
    transforms.RandomAffine(0, shear=10),  # 随机剪切
    transforms.ToTensor(),
])

test_transform = transforms.ToTensor()

# # 重新加载带增强的训练数据
train_ds_aug = torchvision.datasets.MNIST('data', train=True, download=True, transform=train_transform)
# train_dl = torch.utils.data.DataLoader(train_ds_aug, batch_size=32, shuffle=True)
# from torch.utils.data import ConcatDataset
train_ds = torch.utils.data.ConcatDataset([train_ds_original, train_ds_aug])
test_dl  = torch.utils.data.DataLoader(test_ds, 
                                       batch_size=batch_size)
imgs, labels = next(iter(train_dl))
imgs.shape
plt.figure(figsize=(30, 5))
# for i, imgs in enumerate(imgs[:30]):
#     npimg = np.squeeze(imgs.numpy())
#     plt.subplot(3, 10, i + 1)
#     plt.imshow(npimg, cmap=plt.cm.binary)
#     plt.axis('off')
# plt.show()
import torch
import torch.nn as nn
import torch.nn.functional as F
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
num_classes = 10 

class Model(nn.Module):
     def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3) 
        self.pool1 = nn.MaxPool2d(2)                  
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.pool2 = nn.MaxPool2d(2) 
                                      
        self.fc1 = nn.Linear(1600, 64)    

        # self.dropout = nn.Dropout(0.9)   

        self.fc2 = nn.Linear(64, num_classes)
     def forward(self, x):
        x = self.pool1(F.relu(self.conv1(x)))     
        x = self.pool2(F.relu(self.conv2(x)))

        x = torch.flatten(x, start_dim=1)

        x = F.relu(self.fc1(x))

        # x = self.dropout(x)

        x = self.fc2(x)
       
        return x
        
from torchinfo import summary
model = Model().to(device)

summary(model)
#训练
loss_fn    = nn.CrossEntropyLoss() 
learn_rate = 1e-2 
opt        = torch.optim.SGD(model.parameters(),lr=learn_rate)
def train(dataloader, model, loss_fn, optimizer):
    size = len(dataloader.dataset)  # 训练集的大小，一共60000张图片
    num_batches = len(dataloader)  # 批次数目，1875（60000/32）

    train_loss, train_acc = 0, 0  # 初始化训练损失和正确率

    for X, y in dataloader:  # 循环遍历dataloader中的每个批量，每个批量包含一批图片X和对应的标签y。
        X, y = X.to(device), y.to(device)

        # 计算预测误差
        pred = model(X)  # 通过模型对当前批量的图片进行预测
        loss = loss_fn(pred, y)  # 计算预测结果pred和真实标签y之间的损失

        # 反向传播
        optimizer.zero_grad()  # grad属性归零
        loss.backward()  # 反向传播,计算损失关于模型参数的梯度
        optimizer.step()  # 每一步自动更新,根据计算出的梯度更新模型的参数

        # 记录acc与loss
        train_acc += (pred.argmax(1) == y).type(torch.float).sum().item() #更新训练准确率。pred.argmax(1)返回每个预测向量的最大值索引，即最可能的类别。然后比较预测的类别和真实的类别y，计算正确预测的图片数量。
        train_loss += loss.item() #累加当前批量的损失

    train_acc /= size
    train_loss /= num_batches 

    return train_acc, train_loss 
def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)  # 测试集的大小，一共10000张图片
    num_batches = len(dataloader)  # 批次数目，313（10000/32=312.5，向上取整）
    test_loss, test_acc = 0, 0

    with torch.no_grad():
        for imgs, target in dataloader:
            imgs, target = imgs.to(device), target.to(device)

            # 计算loss
            target_pred = model(imgs)
            loss = loss_fn(target_pred, target)

            test_loss += loss.item()
            test_acc += (target_pred.argmax(1) == target).type(torch.float).sum().item()

    test_acc /= size
    test_loss /= num_batches

    return test_acc, test_loss
epochs = 10 # 定义了训练循环的次数，即进行了10次epoch
train_loss = []
train_acc = []
test_loss = []
test_acc = []

for epoch in range(epochs):
    model.train() 
    epoch_train_acc, epoch_train_loss = train(train_dl, model, loss_fn, opt) #train函数计算并返回当前epoch的训练准确率和训练损失

    model.eval() 
    epoch_test_acc, epoch_test_loss = test(test_dl, model, loss_fn) #test函数计算并返回当前epoch的测试准确率和测试损失

    train_acc.append(epoch_train_acc)
    train_loss.append(epoch_train_loss)
    test_acc.append(epoch_test_acc)
    test_loss.append(epoch_test_loss)

    template = ('Epoch:{:2d}, Train_acc:{:.1f}%, Train_loss:{:.3f}, Test_acc:{:.1f}%，Test_loss:{:.3f}')
    print(template.format(epoch + 1, epoch_train_acc * 100, epoch_train_loss, epoch_test_acc * 100, epoch_test_loss))
print('Done')

import matplotlib.pyplot as plt
import warnings
warnings.filterwarnings("ignore")     
plt.rcParams["font.sans-serif"] = ["SimHei"] # 设置字体为黑体
plt.rcParams["axes.unicode_minus"] = False # 正常显示负号
plt.rcParams['figure.dpi']         = 100        #分辨率

epochs_range = range(epochs)

plt.figure(figsize=(12, 3))
plt.subplot(1, 2, 1)

plt.plot(epochs_range, train_acc, label='train accuracy')
plt.plot(epochs_range, test_acc, label='test accuracy')
plt.legend(loc='lower right')
plt.title('train & valid accuracy')

plt.subplot(1, 2, 2)
plt.plot(epochs_range, train_loss, label='train loss')
plt.plot(epochs_range, test_loss, label='test loss')
plt.legend(loc='upper right')
plt.title('train & valid loss')
plt.show()

